import { db } from '../config/database';
import { knowledgeTable, knowledgeInsertSchema } from '../schemas/knowledge';
import { eq, sql } from 'drizzle-orm';

export class KnowledgeRepository {
  private db = db; // 注入或获取 Drizzle 实例

  /**
   * 示例: 创建知识条目 (包含验证)
   * @param data - 未经验证的输入数据
   * @param embedding - 由外部服务 (如 OpenAI) 生成的向量
   */
  async create(data: unknown, embedding: number[]) {
    // 1. 使用 Zod 验证输入数据
    const validatedData = knowledgeInsertSchema.parse(data);

    // 2. 使用 Drizzle 插入数据
    const result = await this.db
      .insert(knowledgeTable)
      .values({
        ...validatedData,
        embedding: embedding,
      })
      .returning({ id: knowledgeTable.id });

    return result[0];
  }

  /**
   * 示例: 向量相似度搜索 (L2 距离)
   * @param embedding - 用于查询的向量
   * @param limit - 返回结果的数量
   */
  async findSimilar(embedding: number[], limit: number = 10) {
    const items = await this.db
      .select({
        id: knowledgeTable.id,
        content: knowledgeTable.content,
        knowledgeType: knowledgeTable.knowledgeType,
        // 使用 pgvector 的 L2 距离操作符 <->
        distance: sql<number>`${knowledgeTable.embedding} <-> ${JSON.stringify(embedding)}::vector`.as('distance'),
      })
      .from(knowledgeTable)
      .orderBy(sql`${knowledgeTable.embedding} <-> ${JSON.stringify(embedding)}::vector`)
      .limit(limit);

    return items;
  }

  /**
   * 示例: 混合查询 (先按元数据过滤，再进行向量搜索)
   * @param embedding - 用于查询的向量
   * @param type - 要过滤的知识类型
   * @param limit - 返回结果的数量
   */
  async findSimilarWithType(embedding: number[], type: 'task' | 'document' | 'conversation', limit: number = 5) {
    const items = await this.db
      .select({
        id: knowledgeTable.id,
        content: knowledgeTable.content,
        // 使用 pgvector 的 L2 距离操作符 <->
        distance: sql<number>`${knowledgeTable.embedding} <-> ${JSON.stringify(embedding)}::vector`.as('distance'),
      })
      .from(knowledgeTable)
      // 1. 先用 WHERE 子句进行精确过滤
      .where(eq(knowledgeTable.knowledgeType, type))
      // 2. 然后在过滤后的结果集上进行排序
      .orderBy(sql`${knowledgeTable.embedding} <-> ${JSON.stringify(embedding)}::vector`)
      .limit(limit);

    return items;
  }
}
